# import numpy as np
# from sklearn.metrics.pairwise import cosine_similarity
#
#
# class PaperRecommender:
#     def __init__(self, embeddings):
#         self.embeddings = embeddings
#
#     def recommend(self, paper_id, top_k=5):
#         target = self.embeddings[paper_id]
#         sim_scores = cosine_similarity([target], self.embeddings)
#         return np.argsort(sim_scores[0])[-top_k - 1:-1]
import numpy as np
import torch
from sklearn.metrics.pairwise import cosine_similarity
from tqdm import tqdm


class PaperRecommender:
    def __init__(self, model, dataset):
        self.model = model
        self.dataset = dataset
        self.embeddings = self._generate_embeddings()

    def _generate_embeddings(self):
        """生成所有论文的特征向量"""
        embeddings = []
        with torch.no_grad():
            for idx in tqdm(range(len(self.dataset))):
                item = self.dataset[idx]
                outputs = self.model(
                    item['input_ids'].unsqueeze(0),
                    item['attention_mask'].unsqueeze(0)
                )
                embeddings.append(outputs.numpy())
        return np.concatenate(embeddings)

    def recommend(self, paper_id, top_k=5):
        target = self.embeddings[paper_id]
        sim_scores = cosine_similarity([target], self.embeddings)
        return np.argsort(sim_scores[0])[-top_k - 1:-1][::-1]